
import torch
import torch.nn as nn
import torch.nn.functional as F
class LSTMModel(nn.Module):

    def __init__(self, vocab_size, embedding_dim, hidden_dim, num_layer):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.num_layer = num_layer
        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        self.lstm = nn.LSTM(embedding_dim, hidden_dim, num_layer)
        self.linear = nn.Linear(hidden_dim, vocab_size)

    def forward(self, input, hidden=None):  # input:[L-1 N]
        seq_len, batch_size = input.size()
        if hidden is None:
            h_0 = input.data.new(2, batch_size, self.hidden_dim).fill_(0).float()
            c_0 = input.data.new(2, batch_size, self.hidden_dim).fill_(0).float()
        else:
            h_0, c_0 = hidden
        embeds = self.embedding(input)  #  embeds：[L-1, N, embedding_dim]
        output, hidden = self.lstm(embeds, (h_0, c_0))  # 所有时间步的输出序列: [L-1, N, hidden_dim] 最终的隐状态和细胞状态: ([2, N, hidden_dim], [2, N, hidden_dim])
        output = self.linear(output)  # [L-1, N, vocab_size]
        return output, hidden

class MaskedSoftmaxCELoss(nn.CrossEntropyLoss):
    def forward(self, pred, target, invalid_idx):
        '''
        pred shape : batch_size , vocab_size , seq_len
        target shape : batch_size , seq_len
        '''
        weights = torch.ones_like(target).fill_(invalid_idx)
        weights = (weights != target).float()
        self.reduction = "none"
        unweighted_loss = super(MaskedSoftmaxCELoss, self).forward(pred, target)
        weighted_loss = (unweighted_loss*weights)
        return weighted_loss.mean()

